import re
import matplotlib.pyplot as plt

def extract_loss_from_log(log_file_path):
    """
    从日志文件中提取 loss 值。
    
    参数:
        log_file_path (str): 日志文件的路径。
    
    返回:
        list: 包含所有提取到的 loss 值的列表。
    """
    loss_values = []
    with open(log_file_path, 'r') as file:
        for line in file:
            # 使用正则表达式匹配 loss 值
            match = re.search(r'loss:([\d.]+)', line)
            if match:
                # 将匹配到的 loss 值转换为浮点数并添加到列表中
                loss_value = float(match.group(1))
                loss_values.append(loss_value)
    return loss_values

def plot_losses(loss_data, save_path):
    """
    绘制多个 loss 值的变化曲线图，并保存图像。
    
    参数:
        loss_data (dict): 包含多个日志文件的损失值的字典，键为日志文件名，值为损失值列表。
        save_path (str): 保存图像的路径。
    """
    # 绘制曲线图
    plt.figure(figsize=(10, 6))
    
    # 为不同的日志文件指定不同的颜色
    colors = ['blue', 'orange', 'green', 'red', 'purple', 'brown', 'pink', 'gray', 'olive', 'cyan']
    
    for i, (log_name, loss_values) in enumerate(loss_data.items()):
        steps = list(range(len(loss_values)))
        plt.plot(steps, loss_values, label=log_name, color=colors[i % len(colors)], linestyle='-', marker='o', markersize=3, alpha=0.7)
    
    plt.title('Loss Over Training Steps')
    plt.xlabel('Step')
    plt.ylabel('Loss')
    plt.legend(loc='upper right', fontsize='small')  # 设置图例位置和字体大小
    plt.grid(True, linestyle='--', alpha=0.5)  # 设置网格线样式
    plt.tight_layout()  # 自动调整子图参数，使之填充整个图像区域
    
    # 保存图像
    plt.savefig(save_path)
    print(f"图像已保存到: {save_path}")
    plt.show()

if __name__ == "__main__":
    # 日志文件路径列表
    log_file_paths = {
        # 添加更多日志文件路径
    }
    # 图像保存路径
    save_path = 'loss_vlm_sft_1B.png'  # 替换为你希望保存图像的路径
    
    # 提取 loss 值
    loss_data = {}
    for log_name, log_file_path in log_file_paths.items():
        loss_data[log_name] = extract_loss_from_log(log_file_path)
    
    # 绘制曲线图并保存
    plot_losses(loss_data, save_path)